#!/usr/bin/env python3
# -*- encoding: utf-8 -*-


import sys
from collections import OrderedDict
sys.path.append('..')
from common.functions import *
from common.gradient import numerical_gradient
from common.layers import *


class TwoLayerNet:

    def __init__(self, input_size, hidden_size, output_size, weight_init_std=0.01, params=None):
        """
        构造函数
        :param input_size: 输入层神经元数目
        :param hidden_size: 中间层神经元数目
        :param output_size: 输出层神经元数目
        :param weight_init_std: 初始化时权重的标准差
        """
        # 初始化权重和偏移
        if params is None:
            self.params = {
                'W1': weight_init_std * np.random.randn(input_size, hidden_size),
                'b1': np.zeros(hidden_size),
                'W2': weight_init_std * np.random.randn(hidden_size, output_size),
                'b2': np.zeros(output_size)
            }
        else:
            self.params = params

        # 构造层
        self.layers = OrderedDict()
        self.layers['Affine1'] = Affine(self.params['W1'], self.params['b1'])
        self.layers['ReLU1'] = ReLU()
        self.layers['Affine2'] = Affine(self.params['W2'], self.params['b2'])

        self.lastlayer = SoftmaxWithLoss()

    # 预测函数
    def predict(self, x):
        for layer in self.layers.values():
            x = layer.forward(x)

        return x

    # 损失函数
    def loss(self, x, t):
        y = self.predict(x)

        return self.lastlayer.forward(y, t)

    # 预测准确度
    def accuracy(self, x, t):
        y = self.predict(x)
        y = np.argmax(y, axis=1)
        t = np.argmax(t, axis=1) if t.ndim != 1 else t

        res = np.sum(y == t) / float(y.shape[0])
        return res

    # x:输入数据, t:监督数据
    def numerical_gradient(self, x, t):
        loss_W = lambda W: self.loss(x, t)

        grads = {
            'W1': numerical_gradient(loss_W, self.params['W1']),
            'b1': numerical_gradient(loss_W, self.params['b1']),
            'W2': numerical_gradient(loss_W, self.params['W2']),
            'b2': numerical_gradient(loss_W, self.params['b2'])
        }

        return grads

    # 使用反向误差传播的方法求梯度
    def gradient(self, x, t):
        # 正向传播
        self.loss(x, t)

        # 反向传播
        dout = 1
        dout = self.lastlayer.backward(dout)

        layers = list(self.layers.values())
        layers.reverse()
        for layer in layers:
            dout = layer.backward(dout)

        grads = {
            'W1': self.layers['Affine1'].dW,
            'b1': self.layers['Affine1'].db,
            'W2': self.layers['Affine2'].dW,
            'b2': self.layers['Affine2'].db
        }

        return grads
